package rpcserver

import (
	"context"
	"fmt"
	"k8s-gpu-collector/proto"
	"net"
	"net/http"
	"os"
	"os/signal"
	"strings"
	"syscall"
	"time"

	"github.com/prometheus/common/expfmt"
	"github.com/sirupsen/logrus"
	"go.uber.org/atomic"
	"golang.org/x/sync/errgroup"
	"google.golang.org/grpc"
	"google.golang.org/protobuf/types/known/emptypb"
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
	"k8s.io/client-go/kubernetes"
)

const (
	ExporterNamespace    = "dcgm"
	DCGMExporterPort     = "9400"  // DCGM exporter默认端口
	DCGMExporterPortTest = "30040" // 测试端口，正式环境下无用
	MetricsPath          = "/metrics"
	RequestTimeout       = 30 * time.Second
)

type RpcServer struct {
	proto.UnimplementedGPUMetricsServiceServer

	cs         *kubernetes.Clientset
	port       int
	httpClient *http.Client
}

func (s *RpcServer) GetOverAllGPUMetrics(ctx context.Context, _ *emptypb.Empty) (*proto.GPUMetricsResponse, error) {
	pods, err := s.cs.CoreV1().Pods(ExporterNamespace).List(ctx, metav1.ListOptions{
		//LabelSelector: "app=dcgm-exporter", // 这里可以根据实际情况调整标签选择器，但是daemonset貌似没有什么可以使用的
	})
	if err != nil {
		return nil, err
	}

	// 用于收集所有指标数据
	var allMetrics []string
	free := atomic.Int32{}
	used := atomic.Int32{}
	reserved := atomic.Int32{}

	// 对每个候选服务都发送请求
	eg := errgroup.Group{}
	for _, pod := range pods.Items {
		if !strings.HasPrefix(pod.Name, "dcgm-exporter") {
			continue
		}
		eg.Go(func() error {
			// 解析Pod IP
			if pod.Status.PodIP == "" {
				logrus.Warnf("Pod %s has no IP address", pod.Name)
				return nil
			}

			// 构建请求URL
			metricsURL := fmt.Sprintf("http://%s:%s%s", pod.Status.PodIP, DCGMExporterPort, MetricsPath)

			// 创建HTTP请求
			req, err := http.NewRequestWithContext(ctx, "GET", metricsURL, nil)
			if err != nil {
				logrus.Errorf("创建请求失败 for pod %s: %v", pod.Name, err)
				return nil // 不返回错误，继续处理其他Pod
			}

			// 发送HTTP请求
			resp, err := s.httpClient.Do(req)
			if err != nil {
				logrus.Errorf("请求DCGM exporter失败 for pod %s (%s): %v", pod.Name, metricsURL, err)
				return nil
			}
			defer resp.Body.Close()

			if resp.StatusCode != http.StatusOK {
				logrus.Errorf("DCGM exporter返回非200状态码 for pod %s: %d", pod.Name, resp.StatusCode)
				return nil
			}

			// 读取响应数据
			//body, err := io.ReadAll(resp.Body)
			//if err != nil {
			//	logrus.Errorf("读取响应数据失败 for pod %s: %v", pod.Name, err)
			//	return nil
			//}

			// metric analyse，提取我们感兴趣的free、used和reserved指标
			parser := expfmt.TextParser{}
			result, err := parser.TextToMetricFamilies(resp.Body)
			if err != nil {
				logrus.Errorf("解析Prometheus指标数据失败 for pod %s: %v", pod.Name, err)
				return nil
			}
			for name, v := range result {
				if name != "DCGM_FI_DEV_FB_FREE" && name != "DCGM_FI_DEV_FB_USED" && name != "DCGM_FI_DEV_FB_RESERVED" {
					continue
				}
				if name == "DCGM_FI_DEV_FB_FREE" {
					for _, m := range v.Metric {
						if m.Gauge == nil || m.Gauge.Value == nil {
							continue
						}
						free.Add(int32(*m.Gauge.Value))
					}
				} else if name == "DCGM_FI_DEV_FB_USED" {
					for _, m := range v.Metric {
						if m.Gauge == nil || m.Gauge.Value == nil {
							continue
						}
						used.Add(int32(*m.Gauge.Value))
					}
				} else if name == "DCGM_FI_DEV_FB_RESERVED" {
					for _, m := range v.Metric {
						if m.Gauge == nil || m.Gauge.Value == nil {
							continue
						}
						reserved.Add(int32(*m.Gauge.Value))
					}
				}
			}
			return nil
		})
	}

	err = eg.Wait()
	if err != nil {
		logrus.Errorf("[GetOverAllGPUMetrics] - 获取GPU指标失败: %v", err)
	}

	// TODO: 这里后续需要解析和处理收集到的Prometheus指标数据
	logrus.Infof("总共收集到 %d 个Pod的GPU指标数据", len(allMetrics))

	// 暂时返回空响应，后续处理指标数据
	return &proto.GPUMetricsResponse{
		Free: float64(free.Load()),
		Used: float64(used.Load()) + float64(reserved.Load()),
	}, nil
}

func NewRpcServer(cs *kubernetes.Clientset, port int) *RpcServer {
	return &RpcServer{
		cs:   cs,
		port: port,
		httpClient: &http.Client{
			Timeout: RequestTimeout,
		},
	}
}

func RunServer(s *RpcServer) error {
	lis, err := net.Listen("tcp", fmt.Sprintf(":%d", s.port))
	if err != nil {
		logrus.Errorf("[RunServer] - failed to listen: %v", err)
		return err
	}
	srv := grpc.NewServer()
	proto.RegisterGPUMetricsServiceServer(srv, s)
	logrus.Infof("[RunServer] - Starting gRPC server on port: %d", s.port)
	go func() {
		if err := srv.Serve(lis); err != nil {
			logrus.Errorf("[RunServer] - failed to serve: %v", err)
			return
		}
	}()
	stopCh := make(chan os.Signal)
	signal.Notify(stopCh, syscall.SIGINT, syscall.SIGTERM)
	<-stopCh

	err = lis.Close()

	if err != nil {
		logrus.Errorf("[RunServer] - failed to close listener: %v", err)
		return err
	}

	return nil
}
